import os
from xml.etree import ElementTree as ET

from shutil import copy
start_path = '/home/jlgpu1/datasets/PCB_DATASET'
new_path = '/home/jlgpu1/datasets/'

annotation_path = start_path + '/Annotations'
images_path = start_path + '/images'
class_labels_path = '/labels'
detect_labels = {'missing_hole': 0, 'mouse_bite': 1, 'open_circuit': 2, 'short': 3, 'spur': 4, 'spurious_copper': 5}
all_labels = {}
new_imgs = '/images'
img_classes = 'img_classes'


def convert_label(tw, th, box):
    dw = 1. / tw
    dh = 1. / th
    x = (box[0] + box[2]) / 2.0
    y = (box[1] + box[3]) / 2.0
    w = box[2] - box[0]
    h = box[3] - box[1]
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return (x, y, w, h)


def xml_yolo():
    file_count = 0
    for item in os.listdir(annotation_path):  # image class label name
        all_labels[item] = {}
        for xml_item in os.listdir(annotation_path + '/' + item):  # label xml file
            file_count += 1
            all_labels[item] = {xml_item: []}
            new_file_name = ''
            # print(xml_item)
            labels = []
            inx = -1
            width = 0.0
            height = 0.0
            etree = ET.parse(annotation_path + '/' + item + '/' + xml_item)
            class_labels = xml_item.split('_')[0]

            new_data_path = new_path + 'pcb_' + class_labels + '/'

            # create labels dir pcb_objclass/labels/class.text
            if not os.path.exists(new_data_path + '/' + 'labels'):
                os.makedirs(new_data_path + '/' + 'labels')

            # create images dir pcb_objclass/images/class.jpg, copy images to new data images path
            if not os.path.exists(new_data_path + '/' + new_imgs):
                os.makedirs(new_data_path + '/' + new_imgs)

            copy(images_path + '/' + item + '/' + xml_item.split('.')
                 [0] + '.jpg', new_data_path + '/' + new_imgs + '/' + str(file_count) + '.jpg')
            for elem in etree.iter():
                if elem.tag == 'width':
                    width = elem.text
                if elem.tag == 'height':
                    height = elem.text
                if elem.tag == 'object':
                    inx = inx + 1
                if elem.tag == 'name':
                    labels.append({'name': elem.text})
                if elem.tag == 'bndbox':
                    labels[inx]['bndbox'] = {}
                    labels[inx]['height'] = height
                    labels[inx]['width'] = width
                if elem.tag == 'xmin':
                    labels[inx]['bndbox']['xmin'] = int(elem.text)
                if elem.tag == 'ymin':
                    labels[inx]['bndbox']['ymin'] = int(elem.text)
                if elem.tag == 'xmax':
                    labels[inx]['bndbox']['xmax'] = int(elem.text)
                if elem.tag == 'ymax':
                    labels[inx]['bndbox']['ymax'] = int(elem.text)

            all_labels[item][xml_item] = labels
            file_num = xml_item.replace('.xml', '').split('_')[-1]
            # if not file_num:
            #     file_num = '00'
            # with open(start_path + class_labels_path + '/' + class_labels + '/' + labels[0]['name'] + file_num + '.txt', 'w') as f:
            with open(new_data_path + '/' + 'labels' + '/' + str(file_count) + '.txt', 'w') as f:
                out_str = ''
                for l_item in labels:
                    out_str += (str(detect_labels[l_item['name']]) + ' ')
                    # covert x，y
                    cov_lbs = convert_label(int(l_item['width']), int(
                        l_item['height']), list(l_item['bndbox'].values()))
                    out_str += (' '.join(map(str, cov_lbs)))

                    out_str += '\n'
                f.write(out_str)

    print(labels)


def trans_imgs(imgs_dir):
    """将图片按照类别抽取到类别文件夹

    Args:
        imgs_dir ([type]): [description]
    """
    all_imgs = {}

    for root, dirs, files in os.walk(imgs_dir, topdown=False):
        for name in files:
            print(os.path.join(root, name))
            img_class = name.split('_')[0]

            if not os.path.exists(new_path + img_classes + '/' + img_class):
                os.makedirs(new_path + img_classes + '/' + img_class)

            copy(os.path.join(root, name), new_path + img_classes + '/' + img_class)

            if img_class not in all_imgs:
                all_imgs[img_class] = []
            all_imgs[img_class].append(os.path.join(root, name))

    print(all_imgs)

    # for name in dirs:
    #     print(os.path.join(root, name))


if __name__ == '__main__':
    xml_yolo()
    # trans_imgs(images_path)
